import io
import json
import os
from PIL import ImageDraw, Image
from concurrent.futures import ThreadPoolExecutor
import threading
import fitz  # PyMuPDF

# 外部模块（根据你环境中已有的）
from milvus_op import MilvusOP
from sentence_transformers import SentenceTransformer as TextEmbeder
from llm_api import VllmModel


def highlight_pdf_region(path, page_number, rects, output_path=None, zoom=2.0):
    """
    在 PDF 指定页面绘制半透明高亮矩形并保存为图片。

    :param path: PDF 文件路径
    :param page_number: 页面编号（从 0 开始）
    :param rects: 区域列表，每个区域是 (x1, y1, x2, y2)
    :param output_path: 输出图片路径
    :param zoom: 放大倍数（控制清晰度）
    """
    try:
        page_number = int(page_number)
        doc = fitz.open(path)
        page = doc.load_page(page_number)

        matrix = fitz.Matrix(zoom, zoom)
        pix = page.get_pixmap(matrix=matrix)
        image = Image.open(io.BytesIO(pix.tobytes("png"))).convert("RGBA")

        overlay = Image.new("RGBA", image.size, (255, 255, 255, 0))
        draw = ImageDraw.Draw(overlay)

        for x1, y1, x2, y2 in rects:
            draw.rectangle(
                (x1 * zoom, y1 * zoom, x2 * zoom, y2 * zoom),
                fill=(255, 0, 0, 80),  # 红色，透明度 80
                outline="red",
                width=2
            )

        combined = Image.alpha_composite(image, overlay)

        if not output_path:
            output_path = f"highlighted_region_{page_number}.png"

        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        combined.save(output_path)
        print(f"[✓] Saved highlighted image: {output_path}")

    except Exception as e:
        print(f"[✗] Failed to highlight page {page_number}: {e}")


def draw_pic_single(item, pdf_path):
    """
    对某一项标注信息进行高亮绘图。
    item 格式: (posInfo, otherData?)
    posInfo 是一个 dict，形如：{page_num: [(x1, y1, x2, y2), ...]}
    """
    posInfo = item[0]
    for page_num, rects in posInfo.items():
        output_path = f"output/retrieval/refer_{page_num}.png"
        highlight_pdf_region(pdf_path, page_num, rects, output_path=output_path)


def draw(img_drawn, pdf_path):
    """
    多线程执行 draw_pic_single
    :param img_drawn: 需要绘图的标注数据列表，每项格式为 (posInfo, ...)
    :param pdf_path: PDF 路径
    """
    os.makedirs("output/retrieval", exist_ok=True)

    with ThreadPoolExecutor(max_workers=4) as executor:
        futures = [executor.submit(draw_pic_single, item, pdf_path) for item in img_drawn]
        for future in futures:
            try:
                future.result()
            except Exception as e:
                print(f"[✗] Error during async drawing: {e}")


# 示例：用于调试或调用入口
if __name__ == "__main__":
    sample_items = [
        (
            {0: [(100, 100, 200, 150), (220, 200, 300, 250)]},
            "extra"
        ),
        (
            {1: [(50, 60, 180, 120)]},
            "extra2"
        )
    ]
    sample_pdf = "example.pdf"
    draw(sample_items, sample_pdf)
